[Dataflow Streaming] [Multi Key] Introduce KeyGroupWorkQueue and integrate with BoundedQueueExecutor by arunpandianp · Pull Request #38767 · apache/beam

arunpandianp · 2026-06-02T09:06:33Z

- Add Uint128Proto and key_group to WorkItem in windmill.proto.
- Update Work model to support key groups.
- Implement KeyGroupWorkQueue for grouping tasks by key group. KepGroupWorkQueue is a FIFO queue that allows polling elements based on global order and also by order within a key group.
- Integrate BoundedQueueExecutor with KepGroupWorkQueue and support targeted polling.

The changes are behind an experiment unstable_enable_multi_key_bundle and does not affect the default core logic.

…oundedQueueExecutor - Add Uint128Proto and key_group to WorkItem in windmill.proto. - Update Work model to support key groups. - Implement KeyGroupWorkQueue for grouping tasks by key group. KepGroupWorkQueue is a FIFO queue that allows polling elements based on global order and also by order within a key group. - Integrate BoundedQueueExecutor with KepGroupWorkQueue and support targeted polling.

arunpandianp · 2026-06-02T09:06:58Z

R: @scwhittle

github-actions · 2026-06-02T09:08:16Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

gemini-code-assist · 2026-06-02T09:10:25Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new KeyGroupWorkQueue implementation to support grouping and targeted polling of work items by key group. It updates the Work model and windmill.proto to include key group information and integrates this new queue mechanism into the BoundedQueueExecutor. These changes are gated behind an experimental flag, ensuring no impact on existing core logic.

Highlights

Protocol Buffers Update: Added Uint128Proto and key_group field to WorkItem in windmill.proto.
Work Model Enhancement: Updated the Work model to support and store key group information.
New Queue Implementation: Introduced KeyGroupWorkQueue, a thread-safe BlockingQueue that supports both global FIFO ordering and targeted polling by computation and key group.
Executor Integration: Integrated KeyGroupWorkQueue into BoundedQueueExecutor, gated behind the unstable_enable_multi_key_bundle experiment.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces support for multi-key bundles in the Streaming Dataflow Worker, guarded by the unstable_enable_multi_key_bundle experiment. It adds a custom, thread-safe KeyGroupWorkQueue to allow polling work by computation and key group in FIFO order, along with corresponding updates to BoundedQueueExecutor, Work, and windmill.proto. Feedback on the new queue implementation suggests adding a defensive check in unlinkNode to prevent double-unlinking issues, and reassigning the result of checkStateNotNull in take() to satisfy static analysis nullness checks.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

arunpandianp · 2026-06-05T08:51:53Z

Run Java PreCommit

scwhittle · 2026-06-05T09:19:34Z

+   */
+  public static final class KeyGroup {
+
+    // The default 0 key group. Work items with 0 keyGroup will always be executed


remove the 0s from the comment , don't think it clarifies anything

scwhittle · 2026-06-05T09:21:29Z

  @GuardedBy("this")
  private long totalTimeMaxActiveThreadsUsed;

+  private final @Nullable KeyGroupWorkQueue keyGroupWorkQueue;


add comment,
// If set the keyGroupWorkQueue is used by the underlying executor.

scwhittle · 2026-06-05T09:33:32Z

+      this.task = task;
+      if (task instanceof QueuedWork) {
+        this.computationId = ((QueuedWork) task).getWork().getComputationId();
+        this.keyGroup = ((QueuedWork) task).getWork().getKeyGroup().orElse(null);


should we leave computationId unset if there is no keygroup? less possible states to worry about.

The comment is on a stale change, now both exist or both does not.

scwhittle · 2026-06-05T09:39:26Z

+        @Nullable Node next = curr.nextNode;
+        unlinkNode(curr);
+        Runnable task = curr.task;
+        if (task != null) {


see above, if task is non-null then we can simplify throughout

done. Updated task to be a nonnull.

scwhittle · 2026-06-05T09:41:59Z

+class KeyGroupWorkQueue extends AbstractQueue<Runnable> implements BlockingQueue<Runnable> {
+
+  static class Node {
+    final @Nullable Runnable task;


instead of nullable how about singleton exception throwing runnable for the sentinels? Seems like it will simplify logic below

Good idea! done.

scwhittle · 2026-06-05T09:43:24Z

+    try {
+      int added = 0;
+      @Nullable Node curr = globalHead.nextNode;
+      while (curr != null && curr != globalTail && added < maxElements) {


throw exception if curr is null? ditto for others

scwhittle

mostly test stuff remaining

scwhittle · 2026-06-05T12:32:11Z


 /** An executor for executing work on windmill items. */
 @SuppressWarnings({
  "nullness" // TODO(https://github.com/apache/beam/issues/20497)


remove? might be able to remove some null checks below then too

scwhittle · 2026-06-05T12:33:08Z

+    if (keyGroupWorkQueue == null) {
+      return null;
+    }
+    QueuedWork queuedWork = keyGroupWorkQueue.pollWork(computationId, keyGroup);


scwhittle · 2026-06-05T12:41:31Z

+    return new QueuedWork(work, executor.createBudgetHandle(1, workBytes));
+  }
+
+  private static class MockRunnable implements Runnable {


nit: NoOpRunnable or FakeRunnable

scwhittle · 2026-06-05T12:43:37Z

+  }
+
+  @Test
+  public void testMemoryPruningLeavesZeroLeaks() {


not sure what memorypruning means here or how leaks is being validated. Can you just add size checking to other tests and remove this one?

scwhittle · 2026-06-05T12:46:24Z

+                    }
+                    if (task != null) {
+                      consumedCount.incrementAndGet();
+                    }


if task is null, should we perform blocking take the next time?

scwhittle · 2026-06-05T12:47:40Z

+    assertTrue(started.await(2, TimeUnit.SECONDS));
+    // Give thread a moment to enter await()
+    Thread.sleep(100);
+    assertEquals(Thread.State.WAITING, t.getState());


instead of sleep and then check, how about a loop while it is not in WAITING

ditto below

scwhittle · 2026-06-05T12:49:49Z

+    t.setDaemon(true);
+    t.start();
+
+    assertTrue(started.await(2, TimeUnit.SECONDS));


higher timeouts here so it isn't flaky? ditto below

scwhittle · 2026-06-05T12:50:38Z

+            () -> {
+              started2.countDown();
+              try {
+                result2.set(queue.poll(2, TimeUnit.SECONDS));


ditto higher so not flaky

scwhittle · 2026-06-05T12:51:08Z

+    assertEquals(workA2, polledA2);
+    assertEquals(1, queue.size());
+
+    // Poll with keyGroup2 again - should return null


trying polling a keygroup never added

github-actions Bot added runners dataflow labels Jun 2, 2026

gemini-code-assist Bot reviewed Jun 2, 2026

View reviewed changes

Comment thread ...ava/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java

Comment thread ...ava/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java Outdated

scwhittle requested changes Jun 4, 2026

View reviewed changes

address comments

5fc07cc

arunpandianp requested a review from scwhittle June 5, 2026 08:20

scwhittle requested changes Jun 5, 2026

View reviewed changes

arunpandianp added 2 commits June 5, 2026 10:25

address comments

b0f8235

Add fairQueue test for KeyGroupWorkQueue

7f7d972

arunpandianp requested a review from scwhittle June 5, 2026 10:34

address comments

4749b08

scwhittle requested changes Jun 5, 2026

View reviewed changes

Conversation

arunpandianp commented Jun 2, 2026

Uh oh!

arunpandianp commented Jun 2, 2026

Uh oh!

github-actions Bot commented Jun 2, 2026

Uh oh!

gemini-code-assist Bot commented Jun 2, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

arunpandianp commented Jun 5, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

scwhittle left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers